//
//  Graph.swift
//  Chapter3LinkedList
//
//  Created by TRS-Mobile on 2021/9/3.
//

import Foundation

/// 边线类型
enum EdgeType {
    /// 有向
    case directed
    /// 无向
    case undirected
}

protocol Graph {
    associatedtype E
    
    /// 创建顶点
    /// - Parameter data: 顶点数据
    /// - Returns: 顶点
    func createVertex(data: E) -> Vertex<E>
    
    /// 添加一条有向边
    /// - Parameters:
    ///   - source: 出发顶点
    ///   - destination: 目的顶点
    ///   - weight: 权重
    func addDirectedEdge(frome source: Vertex<E>,
                         to destination: Vertex<E>,
                         weight: Double?) -> Void
    
    /// 添加一条无向（双向）边
    /// - Parameters:
    ///   - source: 出发顶点
    ///   - detination: 目的顶点
    ///   - weight: 权重
    func addUndirectedEdge(between source: Vertex<E>,
                           and detination: Vertex<E>,
                           weight: Double?) -> Void
    
    /// 在顶点间添加边
    /// - Parameters:
    ///   - edge: 边类型
    ///   - source: 出发顶点
    ///   - destination: 目的顶点
    ///   - weight: 权重
    func add(_ edge: EdgeType,
             from source: Vertex<E>,
             to destination: Vertex<E>,
             weight: Double?) -> Void
    
    /// 返回链接到顶点的所有边
    /// - Parameter source: 顶点
    /// - Returns: 边数组
    func edges(from source: Vertex<E>) -> [Edge<E>]
    
    /// 返回两个顶点间的边权重
    /// - Parameters:
    ///   - source: 原点
    ///   - destination: 目的顶点
    /// - Returns: 权重
    func weight(from source: Vertex<E>,
                to destination: Vertex<E>) -> Double?
}

extension Graph {
    
    func addUndirectedEdge(between source: Vertex<E>,
                           and detination: Vertex<E>,
                           weight: Double?) -> Void {
        addDirectedEdge(frome: source, to: detination, weight: weight)
        addDirectedEdge(frome: detination, to: source, weight: weight)
    }
    
    func add(_ edge: EdgeType,
             from source: Vertex<E>,
             to destination: Vertex<E>,
             weight: Double?) -> Void {
        switch edge {
        case .directed:
            addDirectedEdge(frome: source, to: destination, weight: weight)
        case .undirected:
            addUndirectedEdge(between: source, and: destination, weight: weight)
        }
    }
}

extension Graph {
    
    /// 广度优先搜索
    /// - Parameter source: 开始顶点
    /// - Returns: 搜索顺序路径
    func breadthFirstSearch(from source: Vertex<E>) -> [Vertex<E>] {
        var queue = QueueStack<Vertex<E>>()
        var enqueued: Set<Vertex<E>> = []
        var visited: [Vertex<E>] = []
        
        queue.enqueue(source)
        enqueued.insert(source)
        while let vertex = queue.dequeue() {
            visited.append(vertex)
            let neighborEdges = edges(from: vertex)
            neighborEdges.forEach { edge in
                if !enqueued.contains(edge.destination) {
                    queue.enqueue(edge.destination)
                    enqueued.insert(edge.destination)
                }
            }
        }
        
        return visited
    }
    
    func depthFirstSearch(from source: Vertex<E>) -> [Vertex<E>] {
        var stack: Stack<Vertex<E>> = []
        var pushed: Set<Vertex<E>> = []
        var visited: [Vertex<E>] = []
        
        stack.push(source)
        pushed.insert(source)
        visited.append(source)
        
        outer: while let vertex = stack.peek() {
            let neighbors = edges(from: vertex)
            guard !neighbors.isEmpty else {
                stack.pop()
                continue
            }
            for edge in neighbors {
                if !pushed.contains(edge.destination) {
                    stack.push(edge.destination)
                    pushed.insert(edge.destination)
                    visited.append(edge.destination)
                    continue outer
                }
            }
            stack.pop()
        }
        
        return visited
    }
    
}
